import os
import shutil

# 路径设置
source_train_dir = os.path.join('data', 'NEU-DET', 'train_reorganized')
source_test_dir = os.path.join('data', 'NEU-DET', 'test')
target_train_dir = os.path.join('data', 'NEU-DET', 'train')
target_valid_dir = os.path.join('data', 'NEU-DET', 'valid')

# 备份原始train目录（如果需要）
original_images_dir = os.path.join('data', 'NEU-DET', 'train', 'images')
backup_dir = os.path.join('data', 'NEU-DET', 'train_original_backup')

if os.path.exists(original_images_dir) and not os.path.exists(backup_dir):
    print(f"备份原始图片目录: {original_images_dir} -> {backup_dir}")
    shutil.copytree(original_images_dir, backup_dir)

# 清空原始目录
if os.path.exists(target_train_dir):
    shutil.rmtree(target_train_dir)
if os.path.exists(target_valid_dir):
    shutil.rmtree(target_valid_dir)

# 创建目标目录
os.makedirs(target_train_dir, exist_ok=True)
os.makedirs(target_valid_dir, exist_ok=True)

# 将重组后的训练集移动到原始train目录
print(f"移动训练集: {source_train_dir} -> {target_train_dir}")
for defect_type in os.listdir(source_train_dir):
    source_defect_dir = os.path.join(source_train_dir, defect_type)
    target_defect_dir = os.path.join(target_train_dir, defect_type)
    
    if os.path.isdir(source_defect_dir):
        shutil.copytree(source_defect_dir, target_defect_dir)
        print(f"- 已复制: {defect_type}")

# 将测试集移动到valid目录
print(f"移动测试集: {source_test_dir} -> {target_valid_dir}")
for defect_type in os.listdir(source_test_dir):
    source_defect_dir = os.path.join(source_test_dir, defect_type)
    target_defect_dir = os.path.join(target_valid_dir, defect_type)
    
    if os.path.isdir(source_defect_dir):
        shutil.copytree(source_defect_dir, target_defect_dir)
        print(f"- 已复制: {defect_type}")

# 打印统计信息
print("\n最终数据集统计信息:")
for defect_type in sorted(os.listdir(target_train_dir)):
    train_dir = os.path.join(target_train_dir, defect_type)
    valid_dir = os.path.join(target_valid_dir, defect_type)
    
    train_count = len(os.listdir(train_dir)) if os.path.exists(train_dir) else 0
    valid_count = len(os.listdir(valid_dir)) if os.path.exists(valid_dir) else 0
    
    print(f"- {defect_type}: 训练集 {train_count}张, 验证集 {valid_count}张")

print("\n数据集重组完成!")
print(f"训练集目录: {target_train_dir}")
print(f"验证集目录: {target_valid_dir}")
print(f"原始图片备份目录: {backup_dir}") 